import json
import os.path
import re
import traceback

import requests
import logging
from conn_manager import DIFY_KNOWLEDGE_DATA


# 初始化日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class FileToDify:
    permission = "all_team_members"

    def __init__(self, url_data: str = DIFY_KNOWLEDGE_DATA['local']):

        self.create_knowledge_data = None
        self.top_k = None
        self.description = None
        self.knowledge_name = None
        self.file_path = None
        self.knowledge_list = {}
        self.current_knowledge = {}
        self.total = 0
        self.queried_document = {}  # knowledge_name: {document_data}
        # self.query_knowledge_list()
        self.base_url  = url_data.split('|')[0]
        api_key = url_data.split('|')[1]
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
        self.auth_headers = {
            'Authorization': f'Bearer {api_key}',
        }

    def create_knowledge(self, knowledge_name: str, description: str,
                         # tags_id: list = None,
                         permission: str = "all_team_members",
                         indexing_technique: str = None, reranking_enable: bool = False, reranking_mode: str = None,
                         reranking_provider_name: str = None, reranking_model_name: str = None,
                         embedding_model: str = None, embedding_model_provider: str = None,
                         top_k: int = 2):
        """
        用于创建空知识库
        CURL EX：
        curl --location --request POST 'https://api.dify.ai/v1/datasets' \
            --header 'Authorization: Bearer {api_key}' \
            --header 'Content-Type: application/json' \
            --data-raw '{"name": "name", "permission": "only_me"}'
        :return: 创建成功返回True，否则返回False
        """
        # 检查知识库名称是否已存在
        self.query_knowledge_list()
        if knowledge_name in self.knowledge_list:
            logger.warning(
                f"DIFY - ⚠️ 知识库名称 '{knowledge_name}' 已存在，ID: {self.knowledge_list[knowledge_name]['id']}")
            self.knowledge_name = knowledge_name
            self.description = description
            self.top_k = top_k
            self.permission = permission
            self.current_knowledge = self.knowledge_list[knowledge_name]
            return True

        self.knowledge_name = knowledge_name
        self.description = description
        self.top_k = top_k
        self.permission = permission

        self.create_knowledge_data = {
            "name": self.knowledge_name,
            "description": self.description,
            "permission": permission,
            "indexing_technique": indexing_technique,
            "retrieval_model": {
                "search_method": "semantic_search",
                "reranking_enable": reranking_enable,
                "reranking_mode": reranking_mode,
                "reranking_model": {
                    "reranking_provider_name": reranking_provider_name,
                    "reranking_model_name": reranking_model_name
                },
                "weights": None,
                "top_k": self.top_k,
                "score_threshold_enabled": False,
                "score_threshold": 0,
            },
            # "tags": tags_id,
            "embedding_model": embedding_model,
            "embedding_model_provider": embedding_model_provider,
        }

        # 发送 POST 请求
        response = requests.post(self.base_url, json=self.create_knowledge_data, headers=self.headers)
        self.current_knowledge = response.json()
        # 打印返回的响应内容
        if response.status_code == 200:
            logger.info("DIFY - ✅ Dify 知识库: " + response.json().get('name') + " - ID: " + response.json().get(
                'id') + " -- CREATED.")
            return True
        else:
            logger.error(
                f"DIFY - ❌ Dify 新建知识库{self.knowledge_name} 请求失败: {response.status_code} - {response.text}")
            raise Exception(f"DIFY - ❌ Dify 新建知识库{self.knowledge_name} 请求失败: {response.status_code} - {response.text}")

    @classmethod
    def create_knowledge_instance(cls, knowledge_name: str, description: str,
                         # tags_id: list = None,
                         url_data: str = DIFY_KNOWLEDGE_DATA['local'],
                         permission: str = "all_team_members",
                         indexing_technique: str = None, reranking_enable: bool = False, reranking_mode: str = None,
                         reranking_provider_name: str = None, reranking_model_name: str = None,
                         embedding_model: str = None, embedding_model_provider: str = None,
                         top_k: int = 2):
        """
        创建知识库并返回FileToDify实例（向后兼容方法）
        :return: FileToDify实例
        """
        instance = cls(url_data=url_data)
        success = instance.create_knowledge(
            knowledge_name=knowledge_name,
            description=description,
            permission=permission,
            indexing_technique=indexing_technique,
            reranking_enable=reranking_enable,
            reranking_mode=reranking_mode,
            reranking_provider_name=reranking_provider_name,
            reranking_model_name=reranking_model_name,
            embedding_model=embedding_model,
            embedding_model_provider=embedding_model_provider,
            top_k=top_k
        )
        return instance

    def query_list(self, list_type: str = 'knowledge', knowledge_id: str = None, list_limit: int = 50):
        """
        查询 Dify 知识库或文档列表
        :param url_data:
        :param list_type: 查询类型，'knowledge' 或 'document'
        :param knowledge_id: 知识库ID（仅在查询文档列表时需要）
        :param list_limit: 每页返回的最大数量
        :return: 查询结果列表
        """
        # 根据查询类型构建URL
        if list_type == 'knowledge':
            url = self.base_url
        elif list_type == 'document':
            if not knowledge_id:
                raise ValueError("查询文档列表需要提供 knowledge_id")
            url = f'{self.base_url}/{knowledge_id}/documents'
        else:
            raise ValueError("list_type 必须是 'knowledge' 或 'document'")

        # 清空现有列表
        if list_type == 'knowledge':
            self.knowledge_list = {}
        else:
            # 找到对应的知识库名称
            knowledge_name = None
            for name, data in self.knowledge_list.items():
                if data['id'] == knowledge_id:
                    knowledge_name = name
                    break

            if not knowledge_name:
                logger.warning(f"DIFY - ❌ 未找到ID为 {knowledge_id} 的知识库")
                return []

            # 初始化或清空该知识库的文档列表
            self.queried_document[knowledge_name] = {}

        # 获取第一页数据，同时获取总数
        page = 1
        params = {
            "page": page,
            "limit": list_limit
        }

        response = requests.get(url, headers=self.headers, params=params)
        if response.status_code != 200:
            logger.error(
                f"DIFY - ❌ Failed while query {list_type} list, Page [{page}] with status code: {response.status_code}, Error: {response.text}")
            return []

        response_data = response.json()
        self.total = response_data.get('total', 0)  # 设置total属性

        # 计算需要的总页数
        total_pages = (self.total + list_limit - 1) // list_limit

        # 处理第一页数据
        for item in response_data.get('data', []):
            if list_type == 'knowledge':
                self.knowledge_list[item['name']] = item
            else:
                self.queried_document[knowledge_name][item['id']] = item

        # 如果有多页，继续获取后续页面
        for current_page in range(page + 1, total_pages + 1):
            params["page"] = current_page
            response = requests.get(url, headers=self.headers, params=params)
            if response.status_code != 200:
                logger.error(
                    f"DIFY - ❌ Failed while query {list_type} list, Page [{current_page}] with status code: {response.status_code}, Error: {response.text}")
                continue

            for item in response.json().get('data', []):
                if list_type == 'knowledge':
                    self.knowledge_list[item['name']] = item
                else:
                    self.queried_document[knowledge_name][item['id']] = item

        return response_data.get('data', [])

    def query_knowledge_list(self, list_limit: int = 50):
        """
        查询 Dify 知识库列表
        :param list_limit: 每页返回的最大数量
        :return: 知识库列表
        """
        return self.query_list(list_type='knowledge', list_limit=list_limit)

    def query_document_list(self, knowledge_id: str, list_limit: int = 50):
        """
        查询 Dify 文档列表
        :param knowledge_id: 知识库ID
        :param list_limit: 每页返回的最大数量
        :return: 文档列表
        """
        return self.query_list(list_type='document', knowledge_id=knowledge_id, list_limit=list_limit)

    def check_document_by_name(self, document_name: str):
        """
        检查当前知识库中是否存在指定名称的文档
        :param document_name: 要检查的文档名称
        :return: 如果存在返回 True，否则返回 False
        """
        if not self.current_knowledge:
            logger.warning("DIFY - ⚠️ 当前未选择知识库")
            return False
            
        knowledge_id = self.current_knowledge.get('id')
        if not knowledge_id:
            logger.warning("DIFY - ⚠️ 当前知识库ID无效")
            return False
            
        # 查询当前知识库的文档列表
        self.query_document_list(knowledge_id)
        
        # 获取当前知识库名称
        knowledge_name = self.current_knowledge.get('name')
        if not knowledge_name or knowledge_name not in self.queried_document:
            logger.warning(f"DIFY - ⚠️ 未找到知识库 {knowledge_name} 的文档列表")
            return False
            
        # 检查文档是否存在
        knowledge_name = knowledge_name
        for doc_id, doc_data in self.queried_document[knowledge_name].items():
            if doc_data.get('name') == document_name + '-TechSpec.txt':
                logger.info(f"DIFY - ✅ 在知识库 {knowledge_name} 中找到文档 {document_name}")
                return doc_id
                
        logger.info(f"DIFY - ℹ️ 在知识库 {knowledge_name} 中未找到文档 {document_name}")
        return False

    def del_knowledge(self, knowledge_id: str):
        """
        Delete the knowledge with the knowledge_id
        :param url_data:
        :param knowledge_id:
        :return:
        """
        del_url = self.base_url + '/' + knowledge_id
        response = requests.delete(del_url, headers=self.headers)

        if response.status_code == 204:
            logger.info(f"DIFY - ✅ Dify 知识库 [{knowledge_id}] 删除成功")
        else:
            logger.info(
                f"DIFY - ❌ Failed to delete dataset. Status code: {response.status_code}, Response: {response.text}")

    def del_deprecated_knowledge(self, desc_string_flag: str = 'maintain by Yuxin', deprecated_links: list = [], force: bool = False):
        for knowledge_name, knowledge_data in self.knowledge_list.items():
            try:
                if not knowledge_data.get('description'):
                    continue
                if force:
                    if desc_string_flag in knowledge_data.get('description', ''):
                        self.del_knowledge(knowledge_data['id'])
                        logger.info(f"DIFY - ✅ Dify 知识库 [{knowledge_name}] 删除成功")
                else:
                    url = re.search(r'link【(.*?)】', knowledge_data.get('description'))
                    if knowledge_data.get('description') and desc_string_flag in knowledge_data.get('description', '') \
                            and url in deprecated_links:
                        self.del_knowledge(knowledge_data['id'])
                        logger.info(f"DIFY - ✅ 过期 Dify 知识库 [{knowledge_name}] 删除成功")
            except Exception as e:
                logger.error(f"DIFY - ❌ Failed to delete dataset. Error: {str(e)} - {traceback.format_exc()}")

    def handle_file(self, knowledge_id: str, file_path: str = '../output/AC500-TechSpec.txt',
                   operation: str = 'upload',
                   document_id: str = None,
                   indexing_technique: str = "high_quality",
                   doc_form: str = "text_model",
                   doc_language: str = "English",
                   embedding_model: str = "text-embedding-3-small",
                   embedding_model_provider: str = "openai",
                   search_method: str = "hybrid_search",
                   reranking_enable: bool = True,
                   reranking_mode: str = "reranking_model",
                   reranking_provider_name: str = "cohere",
                   reranking_model_name: str = "rerank-v3.5",
                   keyword_weight: float = 0.3,
                   vector_weight: float = 0.7,
                   top_k: int = 7):
        """
        处理文件上传或更新操作
        :param knowledge_id: 知识库ID
        :param file_path: 文件路径
        :param operation: 操作类型，'upload' 或 'update'
        :param document_id: 文档ID（仅在更新操作时需要）
        :param indexing_technique: 索引技术
        :param doc_form: 文档形式
        :param doc_language: 文档语言
        :param embedding_model: 嵌入模型
        :param embedding_model_provider: 嵌入模型提供商
        :param search_method: 搜索方法
        :param reranking_enable: 是否启用重排序
        :param reranking_mode: 重排序模式
        :param reranking_provider_name: 重排序提供商
        :param reranking_model_name: 重排序模型名称
        :param keyword_weight: 关键词权重
        :param vector_weight: 向量权重
        :param top_k: 返回结果数量
        :return: 响应结果
        """
        # 根据操作类型构建URL
        if operation == 'upload':
            url = f'{self.base_url}/{knowledge_id}/document/create-by-file'
        elif operation == 'update':
            if not document_id:
                raise ValueError("更新操作需要提供 document_id")
            url = f'{self.base_url}/{knowledge_id}/documents/{document_id}/update-by-file'
        else:
            raise ValueError("operation 必须是 'upload' 或 'update'")
        # 根据文件名判断处理方式
        if 'faq' in file_path:
            # FAQ类型文档使用换行符分割，较小的token数
            segmentation_config = {
                "separator": "},",
                "max_tokens": 1000,
                "chunk_overlap": 100
            }
        else:
            # 其他文档使用#号分割，较大的token数
            segmentation_config = {
                "separator": "# ",
                "max_tokens": 3000,
                "chunk_overlap": 500
            }

        # logger.info(f"DIFY - ℹ️ 处理文件上传或更新操作，URL: {url}")
        data = {
            'data': json.dumps({
                "indexing_technique": indexing_technique,
                "process_rule": {
                    "rules": {
                        "pre_processing_rules": [{
                            "id": "remove_extra_spaces",
                            "enabled": True
                        }, {
                            "id": "remove_urls_emails",
                            "enabled": False
                        }],
                        "segmentation": segmentation_config
                    },
                    "mode": "custom"
                },
                # "enabled": enabled,
                "doc_form": doc_form,
                "doc_language": doc_language,
                "retrieval_model": {
                    "search_method": search_method,
                    "reranking_enable": reranking_enable,
                    "reranking_mode": reranking_mode,
                    "reranking_model": {
                        "reranking_provider_name": reranking_provider_name,
                        "reranking_model_name": reranking_model_name
                    },
                    "weights": {
                        "keyword_setting": {
                            "keyword_weight": keyword_weight
                        },
                        "vector_setting": {
                            "vector_weight": vector_weight,
                            "embedding_model_name": embedding_model,
                            "embedding_provider_name": embedding_model_provider
                        }
                    },
                    "top_k": top_k,
                    "score_threshold_enabled": False,
                    "score_threshold": 0
                },
                "embedding_model": embedding_model,
                "embedding_model_provider": embedding_model_provider
            }),
            'type': 'text/plain'
        }

        # 打开文件并发送请求
        if file_path and os.path.exists(file_path):
            with open(file_path, 'rb') as file:
                files = {
                    'file': file
                }
                response = requests.post(url, headers=self.auth_headers, data=data, files=files)
        elif operation == 'update':
            logger.warning("Dify - 没有发现文件，直接更新。")
            response = requests.post(url, headers=self.auth_headers, data=data)
        else:
            logger.error(f"Dify - ❌ 文件不存在: {file_path}, 或者非更新操作, 非法操作！")
            return
            # raise FileNotFoundError(f"文件不存在: {file_path}")

        # 检查响应状态
        if response.status_code == 200:
            logger.info(f"Dify - ✅ 文件{'上传' if operation == 'upload' else '更新'}成功到 Dify 知识库 [{knowledge_id}]")
            return response.json()
        else:
            error_msg = f"Dify - ❌ 文件{'上传' if operation == 'upload' else '更新'}失败. Status code: {response.status_code}"
            try:
                error_data = response.json()
                if 'message' in error_data:
                    error_msg += f", Error: {error_data['message']}"
                if 'code' in error_data:
                    error_msg += f", Code: {error_data['code']}"
            except:
                error_msg += f", Response: {response.text}"
            
            logger.error(error_msg)
            raise Exception(error_msg)

    def upload_file(self, knowledge_id: str, file_path: str = '../output/test.json', **kwargs):
        """
        上传文件到知识库
        :param knowledge_id: 知识库ID
        :param file_path: 文件路径
        :param kwargs: 其他参数，与 handle_file 函数参数相同
        :return: 响应结果
        """
        return self.handle_file(knowledge_id=knowledge_id, file_path=file_path, operation='upload', **kwargs)

    def update_knowledge(self, knowledge_id: str, document_id: str, file_path: str, **kwargs):
        """
        更新知识库中的文档
        :param knowledge_id: 知识库ID
        :param document_id: 文档ID
        :param file_path: 文件路径
        :param kwargs: 其他参数，与 handle_file 函数参数相同
        :return: 响应结果
        """
        return self.handle_file(knowledge_id=knowledge_id, file_path=file_path, operation='update', document_id=document_id, **kwargs)